"""
基于 LangGraph 的报告生成工作流
"""
from typing import Dict, Any, List, TypedDict
from langgraph.graph import StateGraph
from langgraph.checkpoint.memory import MemorySaver
from .state import ReportState, AgentStatus
from .planner_agent import PlannerAgent
from .structure_agent import StructureAgent
from .writer_agent import WriterAgent
from .quality_agent import QualityAgent
from .utils import format_report_content
import logging
import traceback

logger = logging.getLogger(__name__)

# 定义状态类型
class ReportState(TypedDict):
    requirement: str
    task_type: str
    priority: str
    current_step: str
    final_report: str
    quality_score: float
    logs: list
    errors: list

class ReportGenerationWorkflow:
    """报告生成工作流类"""
    
    def __init__(self):
        logger.info("初始化 ReportGenerationWorkflow...")
        self.planner = PlannerAgent()
        self.structure = StructureAgent()
        self.writer = WriterAgent()
        self.quality = QualityAgent()
        
        # 这些属性将由 task_manager 设置
        self.current_task_id = None
        self.task_manager = None
        
        logger.info("ReportGenerationWorkflow 初始化完成")
        
    def generate_report(self, requirement: str, task_type: str = "report_generation", 
                       priority: str = "normal") -> Dict[str, Any]:
        """生成报告的主要方法（改为同步方法）"""
        try:
            logger.info(f"开始生成报告，需求: {requirement}")
            
            # 初始化状态
            state = {
                "requirement": requirement,
                "task_type": task_type,
                "priority": priority,
                "agent_status": {},
                "logs": [],
                "errors": []
            }
            
            # 1. 规划阶段
            logger.info("执行规划阶段...")
            state["agent_status"]["planner"] = "processing"
            
            # 立即更新状态让前端知道规划开始了
            if self.task_manager and self.current_task_id:
                self.task_manager.update_task_partial_status(self.current_task_id, {
                    "agent_status": state["agent_status"],
                    "progress": 5  # 开始规划，进度5%
                })
            
            plan_result = self.planner.execute(state)
            logger.info(f"规划阶段结果: {plan_result}")
            
            if plan_result["success"]:
                state["agent_status"]["planner"] = "completed"
                # 保存优化后的prompt
                if "optimized_prompt" in plan_result:
                    state["optimized_prompt"] = plan_result["optimized_prompt"]
                elif "content" in plan_result:
                    state["optimized_prompt"] = plan_result["content"]
                
                # 保存计划数据
                if "plan" in plan_result:
                    state["plan"] = plan_result["plan"]
                
                # 立即更新任务状态，让前端能够获取到规划结果
                if self.task_manager and self.current_task_id:
                    self.task_manager.update_task_partial_status(self.current_task_id, {
                        "agent_status": state["agent_status"],
                        "optimized_prompt": state.get("optimized_prompt"),
                        "plan": state.get("plan"),
                        "progress": 25  # 规划阶段完成，进度25%
                    })
                
                logger.info(f"规划阶段完成，优化后的prompt: {state.get('optimized_prompt', '未找到')}")
                logger.info(f"生成的计划: {state.get('plan', '未找到')}")
            else:
                state["agent_status"]["planner"] = "failed"
                error_msg = plan_result.get("error", "规划阶段失败")
                state["errors"].append(error_msg)
                return {
                    "success": False,
                    "error": error_msg,
                    "state": state
                }
            
            # 2. 结构化阶段
            logger.info("=== 开始结构化阶段 ===")
            state["agent_status"]["structure"] = "processing"
            
            # 更新状态
            if self.task_manager and self.current_task_id:
                self.task_manager.update_task_partial_status(self.current_task_id, {
                    "agent_status": state["agent_status"],
                    "progress": 30  # 结构化开始，进度30%
                })
            
            structure_result = self.structure.execute(state)
            logger.info(f"结构化阶段结果: {structure_result}")
            
            if not structure_result["success"]:
                state["agent_status"]["structure"] = "failed"
                return self._handle_error("结构化失败", structure_result["error"], state)
            
            state["agent_status"]["structure"] = "completed"
            state["structure"] = structure_result["content"]
            
            # 更新状态
            if self.task_manager and self.current_task_id:
                self.task_manager.update_task_partial_status(self.current_task_id, {
                    "agent_status": state["agent_status"],
                    "progress": 50  # 结构化完成，进度50%
                })
            
            # 3. 写作阶段
            logger.info("=== 开始写作阶段 ===")
            state["agent_status"]["writer"] = "processing"
            
            # 更新状态
            if self.task_manager and self.current_task_id:
                self.task_manager.update_task_partial_status(self.current_task_id, {
                    "agent_status": state["agent_status"],
                    "progress": 55  # 写作开始，进度55%
                })
            
            writer_result = self.writer.execute(state)
            logger.info(f"写作阶段结果: {writer_result}")
            
            if not writer_result["success"]:
                state["agent_status"]["writer"] = "failed"
                return self._handle_error("写作失败", writer_result["error"], state)
            
            state["agent_status"]["writer"] = "completed"
            state["content"] = writer_result["content"]  # 保存写作结果到 content 字段
            
            # 更新状态
            if self.task_manager and self.current_task_id:
                self.task_manager.update_task_partial_status(self.current_task_id, {
                    "agent_status": state["agent_status"],
                    "content": state["content"],
                    "progress": 75  # 写作完成，进度75%
                })
            
            # 4. 质量检查阶段
            logger.info("=== 开始质量检查阶段 ===")
            state["agent_status"]["quality"] = "processing"
            
            # 更新状态
            if self.task_manager and self.current_task_id:
                self.task_manager.update_task_partial_status(self.current_task_id, {
                    "agent_status": state["agent_status"],
                    "progress": 80  # 质量检查开始，进度80%
                })
            
            quality_result = self.quality.execute(state)
            logger.info(f"质量检查阶段结果: {quality_result}")
            
            if not quality_result["success"]:
                state["agent_status"]["quality"] = "failed"
                return self._handle_error("质量检查失败", quality_result["error"], state)
            
            state["agent_status"]["quality"] = "completed"
            
            # 确保有quality_score
            quality_score = 8.0  # 默认评分
            if "evaluation" in quality_result:
                quality_score = quality_result["evaluation"].get("overall_score", 8.0)
            elif "quality_score" in quality_result:
                quality_score = quality_result["quality_score"]
            
            # 最终更新状态
            if self.task_manager and self.current_task_id:
                self.task_manager.update_task_partial_status(self.current_task_id, {
                    "agent_status": state["agent_status"],
                    "quality_score": quality_score,
                    "progress": 100  # 全部完成，进度100%
                })
            
            # 返回最终结果
            final_result = {
                "success": True,
                "content": state["content"],  # 文章内容
                "optimized_prompt": state.get("optimized_prompt", ""),  # 优化后的prompt
                "plan": state.get("plan", ""),  # 计划数据
                "quality_score": quality_score,
                "agent_status": state["agent_status"],
                "logs": state["logs"],
                "errors": state["errors"]
            }
            
            logger.info(f"报告生成完成，最终结果: {final_result}")
            return final_result
            
        except Exception as e:
            logger.error(f"报告生成异常: {str(e)}")
            logger.error(f"异常堆栈: {traceback.format_exc()}")
            return {
                "success": False,
                "error": f"生成异常: {str(e)}",
                "state": state if 'state' in locals() else {}
            }
    
    def _handle_error(self, stage: str, error: str, state: Dict[str, Any]) -> Dict[str, Any]:
        """处理错误"""
        error_msg = f"{stage}: {error}"
        logger.error(f"处理错误: {error_msg}")
        logger.error(f"错误发生时的状态: {state}")
        
        return {
            "success": False,
            "error": error_msg,
            "content": "",
            "quality_score": 0.0,
            "agent_status": state["agent_status"],
            "logs": state["logs"],
            "errors": state["errors"] + [error_msg]
        }
    
    def _get_agent_status(self) -> Dict[str, str]:
        """获取所有代理的状态"""
        return {
            "planner": "completed",
            "structure": "completed",
            "writer": "completed",
            "quality": "completed"
        }
    
    def get_workflow_status(self, thread_id: str = "report_generation") -> Dict[str, Any]:
        """获取工作流状态"""
        try:
            # 这里可以返回一个简单的状态
            return {
                "progress": 0.0,
                "current_agent": "not_started",
                "agent_status": {},
                "logs": [],
                "has_errors": False
            }
                
        except Exception as e:
            logger.error(f"获取工作流状态失败: {str(e)}")
            return {
                "progress": 0.0,
                "current_agent": "error",
                "agent_status": {},
                "logs": [f"状态查询失败: {str(e)}"],
                "has_errors": True
            } 